import os
import random

import tensorflow.compat.v1 as tf


def split_train(from_data_dir, to_dir, task="mnli", ratio=0.5):
    os.makedirs(to_dir, exist_ok=True)

    train_file = os.path.join(from_data_dir, "train.tsv")

    dev_file = None
    dev_file_matched = None
    dev_file_mismatched = None

    if task != "mnli":
        dev_file = os.path.join(from_data_dir, "dev.tsv")
    else:
        dev_file_matched = os.path.join(from_data_dir, "dev_matched.tsv")
        dev_file_mismatched = os.path.join(from_data_dir, "dev_mismatched.tsv")

    list_samples = []
    header = None
    with open(train_file, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            line = line.strip()

            if task in ["mnli", "mrpc", "qnli", "qqp", "rte", "snli", "sst-2", "wnli"]:
                if i == 0:
                    header = line
                    continue

            if not line:
                continue

            list_samples.append(line)

    random.shuffle(list_samples)

    train_samples_new = list_samples[: int(len(list_samples) * ratio)]
    dev_samples_new = list_samples[int(len(list_samples) * ratio): ]

    to_train_file = os.path.join(to_dir, "train.tsv")
    to_dev_file = os.path.join(to_dir, "dev.tsv")

    to_test_file = None
    to_test_file_matched = None
    to_test_file_mismatched = None

    if task != "mnli":
        to_test_file = os.path.join(to_dir, "test.tsv")
    else:
        to_test_file_matched = os.path.join(to_dir, "test_matched.tsv")
        to_test_file_mismatched = os.path.join(to_dir, "test_mismatched.tsv")


    with open(to_train_file, "w", encoding="utf-8") as f:
        if header:
            f.write(header + "\n")
        for line in train_samples_new:
            f.write(line + "\n")

    with open(to_dev_file, "w", encoding="utf-8") as f:
        if header:
            f.write(header + "\n")
        for line in dev_samples_new:
            f.write(line + "\n")

    if task != "mnli":
        tf.gfile.Copy(
            dev_file,
            to_test_file,
            overwrite=True
        )

    else:
        tf.gfile.Copy(
            dev_file_matched,
            to_test_file_matched,
            overwrite=True
        )
        tf.gfile.Copy(
            dev_file_mismatched,
            to_test_file_mismatched,
            overwrite=True
        )


if __name__ == "__main__":

    # cola
    from_data_dir = "datasets/glue/CoLA"
    to_dir = "datasets/glue_split_train/CoLA"
    split_train(
        from_data_dir, to_dir, task="cola", ratio=0.7
    )

    # mnli
    from_data_dir = "datasets/glue/MNLI"
    to_dir = "datasets/glue_split_train/MNLI"
    split_train(
        from_data_dir, to_dir, task="mnli", ratio=0.95
    )

    # MRPC
    from_data_dir = "datasets/glue/MRPC"
    to_dir = "datasets/glue_split_train/MRPC"
    split_train(
        from_data_dir, to_dir, task="mrpc", ratio=0.7
    )

    # QNLI
    from_data_dir = "datasets/glue/QNLI"
    to_dir = "datasets/glue_split_train/QNLI"
    split_train(
        from_data_dir, to_dir, task="qnli", ratio=0.9
    )

    # QQP
    from_data_dir = "datasets/glue/QQP"
    to_dir = "datasets/glue_split_train/QQP"
    split_train(
        from_data_dir, to_dir, task="qqp", ratio=0.9
    )

    # RTE
    from_data_dir = "datasets/glue/RTE"
    to_dir = "datasets/glue_split_train/RTE"
    split_train(
        from_data_dir, to_dir, task="rte", ratio=0.7
    )

    # SNLI
    from_data_dir = "datasets/glue/SNLI"
    to_dir = "datasets/glue_split_train/SNLI"
    split_train(
        from_data_dir, to_dir, task="snli", ratio=0.7
    )

    # SST-2
    from_data_dir = "datasets/glue/SST-2"
    to_dir = "datasets/glue_split_train/SST-2"
    split_train(
        from_data_dir, to_dir, task="sst-2", ratio=0.9
    )

    # WNLI
    from_data_dir = "datasets/glue/WNLI"
    to_dir = "datasets/glue_split_train/WNLI"
    split_train(
        from_data_dir, to_dir, task="wnli", ratio=0.7
    )
